import os
os.environ["CUDA_VISIBLE_DEVICES"] = "1"
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader
from model.attack import Attack
import datetime
import csv
import kornia

from options import HiDDenConfiguration
from model.encoder_decoder import EncoderDecoder
from noise_layers.jpeg import Jpeg, JpegSS
from noise_layers.gaussian_noise import Gaussian_Noise
from noise_layers.quantization import Quantization
from noise_layers.salt_pepper_noise import SP
from noise_layers.colorjitter import ColorJitter
from modules.UnetPlusPlus_L6 import UNetPlusPlusL6_Attack
import numpy as np
if __name__ == '__main__':
    attack_model = Attack(channels=64).cuda()
    attack_model.load_state_dict(torch.load("/home/zzc/simulator_noisy/Unet-Attack/attack_runs-Q-90/cnn_2025-05-12_18-19-50/model/best_model.pth"))
    attack_model.eval()
    jpeg = Jpeg(Q=90).cuda()
    L1 = nn.L1Loss().cuda()
    L2 = nn.MSELoss().cuda()

    messages_length = 128
    hidden_config = HiDDenConfiguration(H=128, W=128,
                                    message_length=messages_length,
                                    encoder_blocks=4, encoder_channels=64,
                                    decoder_blocks=7, decoder_channels=64,
                                    use_discriminator=True,
                                    use_vgg=False,
                                    discriminator_blocks=3, discriminator_channels=64,
                                    decoder_loss=1,
                                    encoder_loss=0.7,
                                    adversarial_loss=1e-3
                                    )

    encoder_decoder = EncoderDecoder(hidden_config).cuda()
    encoder_decoder.load_state_dict(torch.load("/home/zzc/simulator_noisy/run_base_2025-04-10_14-11-47/model/300_model.pth"))
    encoder_decoder.train()

    train_data_dir = "/home/zzc/simulator_noisy/data/train"  # 训练数据
    test_data_dir = "/home/zzc/simulator_noisy/data/val"     # 测试数据
    batch_size = 32
    log_file = "log.csv"

    image_size = 128
    # 定义图像预处理的转换操作
    data_transforms = {
        'train': transforms.Compose([
            transforms.RandomCrop((image_size, image_size), pad_if_needed=True),
            transforms.ToTensor()
        ]),
        'test': transforms.Compose([
            transforms.CenterCrop((image_size, image_size)),
            transforms.ToTensor()
        ])
    }

    # 创建ImageFolder数据集实例
    train_dataset = ImageFolder(train_data_dir, transform=data_transforms['train'])
    test_dataset = ImageFolder(test_data_dir, transform=data_transforms['test'])

    # 创建数据加载器
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    total_L1loss = 0.0
    total_L2loss = 0.0
    with torch.no_grad():
        for images, _ in train_loader:
            images = images.cuda()
            messages = torch.Tensor(np.random.choice([0, 1], size=[images.size(0), 1, messages_length, messages_length])).cuda()
            stego = encoder_decoder.encoder(images, messages)
            L1loss = L1(attack_model(stego), jpeg(stego))
            L2loss = L2(attack_model(stego), jpeg(stego))
            total_L1loss += L1loss
            total_L2loss += L2loss
            print(L1loss)
            print(L2loss)
   
    print(total_L1loss/len(train_loader))
    print(total_L2loss/len(train_loader))